import cv2
import numpy as np
import time
import matplotlib.pyplot as plt

# 读取灰度图像
img = cv2.imread('demo.jpg', cv2.IMREAD_GRAYSCALE)

# 计算直方图
hist = np.zeros((256, ))
for i in range(img.shape[0]):
    for j in range(img.shape[1]):
        hist[img[i, j]] += 1

# 显示直方图
plt.title("Grayscale histogram")  # 灰度直方图
plt.xlabel("Grayscale value")   # 灰度值
plt.ylabel("Number of pixels")  # 像素数
plt.plot(hist)
plt.show()

import random

def salt_and_pepper_noise(img, prob):
    """
    在图像中添加椒盐噪声
    :param img: 原始图像
    :param prob: 噪声点的概率
    :return: 添加噪声后的图像
    """
    height, width = img.shape
    noise_img = img.copy()
    for i in range(height):
        for j in range(width):
            if random.random() < prob:
                noise_img[i, j] = 0 if random.random() < 0.5 else 255
    return noise_img


def gaosi_noise_img(img):
    """
    在图像中加入高斯噪声
    : param img: 原始图像
    : return 添加噪声后的图像
    """
    # 生成高斯噪声矩阵
    mean = random.uniform(-10, 10)
    stddev = random.uniform(5, 15)
    gaosi_noise = np.zeros(img.shape, dtype=np.uint8)

    for i in range(img.shape[0]):
        for j in range(img.shape[1]):

            u1 = random.random()
            u2 = random.random()
            z1 = np.sqrt(-2 * np.log(u1)) * np.cos(2 * np.pi * u2)
            # z2 = np.sqrt(-2 * np.log(u1)) * np.sin(2 * np.pi * u2)
            gaosi_noise[i, j] = np.clip(img[i, j] + mean + stddev * z1, 0, 255)

    return gaosi_noise

# 添加椒盐噪声
noise_start_time = time.time()
noise_img = salt_and_pepper_noise(img, 0.2)
noise_end_time = time.time()

gaosi_start_time = time.time()
gaosi_img = gaosi_noise_img(img) + img
gaosi_end_time = time.time()

# 显示结果
cv2.imwrite('demo_noise.jpg', noise_img)
cv2.imwrite('demo_gaosi_img.jpg', gaosi_img)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

# 计算代码执行的时间并输出结果
time1 = noise_end_time - noise_start_time
time2 = gaosi_end_time - gaosi_start_time
print("添加椒盐噪声执行的时间为：{:.3f}秒".format(time1))
print("添加高斯噪声执行的时间为：{:.3f}秒".format(time2))

# 计算信噪比
MAX = 255
MSE = ((img - noise_img) ** 2).sum() / (img.shape[0] * img.shape[1])

if MSE == 0:
    PSNR = 100
else:
    PSNR = 10 * np.log10(MAX**2 / MSE)


MSE_gaosi = ((img - gaosi_img) ** 2).sum() / (img.shape[0] * img.shape[1])
if MSE_gaosi == 0:
    PSNR_gaosi = 100
else:
    PSNR_gaosi = 10 * np.log10(MAX ** 2 / MSE_gaosi)

print('椒盐噪声的信噪比为：', PSNR)
print('高斯噪声的信噪比为：', PSNR_gaosi)